import requests
import jieba
import pdfplumber
import os
import faiss
import numpy as np


def ask_glm(content, max_length=1000, retries=3):
    # 向GLM模型发送生成请求的函数
    url = os.getenv("GLM_API_URL", "http://127.0.0.1:11434/api/generate")  # 获取GLM API的URL
    data = {
        "model": "qwen2:0.5b",  # 指定要使用的GLM模型
        "prompt": content,  # 设置生成的提示内容
        "max_tokens": max_length,  # 设置生成的最大长度
        "stream": False  # 控制是否使用流式传输
    }
    for attempt in range(retries):  # 尝试多次发送请求
        try:
            response = requests.post(url, json=data)  # 发送POST请求
            if response.status_code == 200:  # 如果响应状态码为200（表示成功）
                response_json = response.json()  # 解析响应的JSON数据
                print(f"API响应：{response_json}")  # 输出调试信息
                return response_json  # 返回生成的响应JSON
            else:  # 如果响应状态码不是200
                print(f"请求失败，状态码：{response.status_code}, 响应内容：{response.text}")  # 输出错误信息
        except Exception as e:  # 处理可能的异常
            print(f"请求大模型API时发生错误：{e}")  # 输出错误信息
        print(f"重试 {attempt + 1}/{retries}...")  # 输出重试信息
    return None  # 如果多次重试仍然失败，则返回None


def load_pdf_content(file_path):
    # 加载PDF文件内容的函数
    try:
        with pdfplumber.open(file_path) as pdf:  # 使用pdfplumber打开PDF文件
            pdf_content = [{'page': i, 'content': page.extract_text()} for i, page in
                           enumerate(pdf.pages)]  # 提取每一页的文本内容
        return pdf_content  # 返回包含页码和内容的列表
    except Exception as e:  # 处理可能的异常
        print(f"读取PDF文件时发生错误：{e}")  # 输出错误信息
        return []  # 如果出现异常，则返回空列表


def preprocess_pdf_content(pdf_content):
    # 预处理PDF内容的函数，使用分词工具进行分词
    pdf_content_words = [jieba.lcut(page['content']) for page in pdf_content]  # 对每一页的内容进行分词处理
    return pdf_content_words  # 返回分词后的内容列表


def build_faiss_index(pdf_content_words):
    # 将每个页面的分词结果转换为向量
    vectors = [np.array([hash(word) % 1000000 for word in words], dtype=np.float32) for words in pdf_content_words]

    # 找到最长向量的长度
    max_len = max(len(vec) for vec in vectors)

    # 对所有向量进行填充，使其长度一致
    padded_vectors = np.array([np.pad(vec, (0, max_len - len(vec))) for vec in vectors], dtype=np.float32)

    # 创建 Faiss 索引
    index = faiss.IndexFlatL2(padded_vectors.shape[1])
    index.add(padded_vectors)
    return index, max_len


def retrieve_top_k_pages_faiss(query, index, pdf_content_words, max_len, k=10):
    query_words = jieba.lcut(query)
    query_vector = np.array([hash(word) % 1000000 for word in query_words], dtype=np.float32)

    # 对查询向量进行填充，使其长度与索引向量一致
    query_vector = np.pad(query_vector, (0, max_len - len(query_vector)))
    query_vector = query_vector.reshape(1, -1)

    _, top_k_indices = index.search(query_vector, k)
    return ['page_' + str(idx + 1) for idx in top_k_indices[0]]


def main():
    pdf_file_path = "美军联合CAS手册2014.pdf"  # PDF文件路径
    pdf_content = load_pdf_content(pdf_file_path)  # 加载PDF文件内容
    if not pdf_content:  # 如果加载失败
        print("无法加载PDF内容，请检查文件路径和文件格式。")  # 输出错误信息
        return  # 终止程序

    pdf_content_words = preprocess_pdf_content(pdf_content)  # 预处理PDF内容，进行分词处理
    faiss_index, max_len = build_faiss_index(pdf_content_words)  # 构建 Faiss 索引

    while True:  # 进入交互循环
        user_question = input("请输入你的问题（输入'q'退出）：")  # 获取用户输入的问题
        if user_question.lower() == 'q':  # 如果用户输入'q'，则退出循环
            break

        faiss_references = retrieve_top_k_pages_faiss(user_question, faiss_index, pdf_content_words,
                                                      max_len)  # 使用 Faiss 检索前 k 个相关页面的页面号

        reference_content = ''
        for ref in faiss_references:
            page_num = int(ref.split('_')[1])  # 获取页面号
            reference_content += pdf_content[page_num - 1]['content'].replace('\n', ' ') + ' '  # 获取页面内容并拼接

        # 构建模型生成的提示内容
        prompt = '''你是一个飞行知识专家，帮我结合给定的资料，用中文回答问题。如果问题无法从资料中获得，请输出结合给定的资料，无法回答问题。
资料：{0}

问题：{1}
        '''.format(reference_content, user_question)

        # 调用 GLM 模型生成答案
        answer = ask_glm(prompt, 1000)
        if answer:
            message_content = answer.get('response', '')  # 获取生成的回答内容
            print(f"message_content内容：{message_content}")  # 输出调试信息
            if '无法' in message_content or not message_content:  # 如果回答内容表明无法回答问题
                message_content = '结合给定资料，无法回答问题'  # 更新回答内容
        else:
            message_content = '无法获取答案，请稍后重试。'  # 如果无法获取答案，则输出错误信息

        # 输出用户问题、参考资料和模型回答
        print(f"【用户提问】\n{user_question}")
        print(f"【参考资料】\n{reference_content}")
        print(f"【模型回答】\n{message_content}")
        print('\n\n\n')


if __name__ == "__main__":
    main()  # 调用主函数，启动程序